import cv2
import mediapipe as mp

def MC_run():
    InputFileName = 'girl'
    InputFileFormat = '.mp4'
    input_path = InputFileName + InputFileFormat

    output_path = 'Annotated_' + input_path

    """获取帧数和声明用于储存关键点信息的数组变量"""
    # 将视频帧数赋值给frame_count变量。
    cap = cv2.VideoCapture(input_path)
    frame_count = 0
    while (cap.isOpened()):
        success, frame = cap.read()
        frame_count += 1
        if not success:
            break
    cap.release()

    # 声明数组变量，赋予空值。
    PoseWorldLandmarkList = []  # 世界空间人体姿态关键点坐标数组。
    PoseWorldLandmarkListOnNthFrame = []  # 每一帧的世界空间人体姿态关键点坐标数组。
    PoseLandmarkList = []  # 屏幕空间人体姿态关键点坐标数组。
    PoseLandmarkListOnNthFrame = []  # 每一帧的屏幕空间人体姿态关键点坐标数组。
    LeftHandLandmarkList = []  # 左手关键点坐标数组。
    LeftHandLandmarkListOnNthFrame = []  # 每一帧的左手关键点坐标数组。
    RightHandLandmarkList = []  # 右手关键点坐标数组。
    RightHandLandmarkListOnNthFrame = []  # 每一帧的右手关键点坐标数组。

    '''
    注：
    1. 关键点坐标数组为三维数组，第一维是帧号，第二维是关键点ID，第三维是xyz坐标。如：PoseLandmarkList[0][0][0]表示
    第一帧第一个关键点（编号为0的关键点）的x坐标；
    2. 每一帧的关键点坐标数组为二维数组，第一维是关键点ID，第二维是xyz坐标。如：PoseLandmarkListOnNthFrame[0][0]
    表示第一个关键点的x坐标；
    '''

    """存储关键点信息，并生成标注了关键点的视频文件"""
    # 获取屏幕的宽、高。
    cap = cv2.VideoCapture(input_path)
    frame_size = (cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    h = frame_size[1]
    w = frame_size[0]

    # 尺寸因子。由于Mediapipe所描述的人体的尺寸与3ds max场景中的人物尺寸存在差异，此处需要声明一个尺寸因子。
    SizeFactor = 140

    # 设置视频输出参数。
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    fps = cap.get(cv2.CAP_PROP_FPS)
    out = cv2.VideoWriter(output_path, fourcc, fps, (int(frame_size[0]), int(frame_size[1])))

    # 开始实施Mediapipe人体姿态关键点识别。
    with mp.solutions.holistic.Holistic(
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5) as holistic:
        while cap.isOpened():
            success, image = cap.read()
            if not success:
                break
            image = cv2.cvtColor(cv2.flip(image, 1), cv2.COLOR_BGR2RGB)
            image.flags.writeable = False
            results = holistic.process(image)
            image.flags.writeable = True
            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
            annotated_image = image.copy()

            # 人体姿态
            if results.pose_landmarks:
                # 如果检测出关键点，则令这一帧的关键点坐标数组清零。
                PoseLandmarkListOnNthFrame = []
                PoseWorldLandmarkListOnNthFrame = []
                for i in [11, 12, 13, 14, 15, 16, 23, 24, 25, 26, 27, 28, 31, 32]:  # 依次为：肩 肘 手腕 髋关节 膝盖 踝关节 脚尖。

                    # 屏幕空间
                    px = results.pose_landmarks.landmark[i].x * w
                    py = results.pose_landmarks.landmark[i].y * h
                    pz = results.pose_landmarks.landmark[i].z * w  # 根据官方说明文件，z坐标需要乘屏幕宽度。

                    # 世界空间
                    cx = results.pose_world_landmarks.landmark[i].x * SizeFactor
                    cy = -results.pose_world_landmarks.landmark[i].y * SizeFactor
                    cz = results.pose_world_landmarks.landmark[i].z * SizeFactor  # 需要乘尺寸因子。

                    # 将各个关键点坐标依次存入数组。
                    PoseWorldLandmarkListOnNthFrame.append([cx, cy, cz])
                    PoseLandmarkListOnNthFrame.append([px, h - py, pz])

                    # 给视频文件的每一帧绘制人体姿态关键点标记，此例为蓝色实心圆。
                    annotated_image = cv2.circle(annotated_image, (round(px), round(py)), 10, (255, 144, 30), -1)
            else:
                # 如果没有检测出关键点，则这一帧的关键点坐标数组以字符串'NaN'标记。在Maxscript中，每当遇到'NaN'，则不创建关键帧。
                PoseLandmarkListOnNthFrame = 'NaN'
                PoseWorldLandmarkListOnNthFrame = 'NaN'

            # 将每一帧的关键点坐标数组作为子元素存入总的关键点坐标数组，从而实现总数组的第一维是帧号，第二维是每一帧的关键点坐标数组。
            PoseLandmarkList.append(PoseLandmarkListOnNthFrame)
            PoseWorldLandmarkList.append(PoseWorldLandmarkListOnNthFrame)

            # 左手
            if results.left_hand_landmarks:
                LeftHandLandmarkListOnNthFrame = []
                for i in range(1, 21):
                    lx = results.left_hand_landmarks.landmark[i].x * w
                    ly = results.left_hand_landmarks.landmark[i].y * h
                    lz = results.left_hand_landmarks.landmark[i].z * w + (
                                results.pose_landmarks.landmark[15].z - results.left_hand_landmarks.landmark[0].z) * w
                    # 手部关键点的坐标原点与人体姿态关键点的坐标原点不同，因而以上语句的赋值符号右边的第二项
                    # (results.pose_landmarks.landmark[15].z-results.left_hand_landmarks.landmark[0].z)*w，其作用
                    # 是使手部的手腕关键点与人体姿态的手腕关键点重合，从而人体姿态和手部能衔接自然。

                    LeftHandLandmarkListOnNthFrame.append([lx, h - ly, lz])

                    # 给视频文件的每一帧绘制手部关键点标记，此例为绿色实心圆。
                    annotated_image = cv2.circle(annotated_image, (round(lx), round(ly)), 5, (0, 255, 0), -1)
            else:
                LeftHandLandmarkListOnNthFrame = 'NaN'

            LeftHandLandmarkList.append(LeftHandLandmarkListOnNthFrame)

            # 右手
            if results.right_hand_landmarks:
                RightHandLandmarkListOnNthFrame = []
                for i in range(1, 21):
                    rx = results.right_hand_landmarks.landmark[i].x * w
                    ry = results.right_hand_landmarks.landmark[i].y * h
                    rz = results.right_hand_landmarks.landmark[i].z * w + (
                                results.pose_landmarks.landmark[16].z - results.right_hand_landmarks.landmark[0].z) * w

                    RightHandLandmarkListOnNthFrame.append([rx, h - ry, rz])

                    annotated_image = cv2.circle(annotated_image, (round(rx), round(ry)), 5, (0, 255, 0), -1)
            else:
                RightHandLandmarkListOnNthFrame = 'NaN'

            RightHandLandmarkList.append(RightHandLandmarkListOnNthFrame)

            if success == True:
                out.write(annotated_image)

    cv2.destroyAllWindows()
    out.release()
    cap.release()

    """采用屏幕空间还是世界空间"""
    CharacterPose = PoseWorldLandmarkList
    # CharacterPose = PoseLandmarkList

    """生成Maxscript脚本，制作火柴人模型的关键帧动画"""
    FileName = 'SetMotion_' + InputFileName + '.ms'
    file = open(FileName, 'w')  # 开始书写Maxscript脚本
    file.write('animate on\n')  # 开始创建关键帧
    file.write('(\n')

    '''
    注：
    3ds Max中的人物的动作由骨骼驱动，骨骼的动作又由一系列控制点的位置变化驱动。控制点作为3ds Max中的对象，可以用
    Maxscript脚本来控制其位置。

    将动作捕捉数据输入3ds Max以制作成动画，本质上是令3ds Max中的每个控制点在每一帧的位置坐标等于Mediapipe
    给出的对应的关键点坐标。

    此例中，Mediapipe给出的第j个人体姿态关键点对应3ds Max场景中的第j个控制点Cube_j；第j个左手关键点对应3ds Max
    场景中的第j个控制点L_j；第j个右手关键点对应3ds Max场景中的第j个控制点R_j。
    '''

    for i in range(frame_count - 1):
        # 写入人体姿态关键帧
        if isinstance(CharacterPose[i], list):  # 如果第i帧关键点数据不是'NaN'，即第i帧含有关键点数据，则创建关键帧。
            for j in range(14):
                file.write('\t' + 'at time ' + str(i) + ' $Cube_' + str(j) + '.pos = [' + str(
                    CharacterPose[i][j][0]) + ', ' + str(CharacterPose[i][j][2]) + ', ' + str(
                    CharacterPose[i][j][1]) + ']\n')
                # 在第i帧，对象Cube_j的位置坐标为[CharacterPose[i][j][0], CharacterPose[i][j][1], CharacterPose[i][j][2]]

        '''
        注：手部关键帧只有屏幕空间，不能与世界空间人体姿态关键点相容。如果采用世界空间人体姿态关键点，则
        请跳过以下关于手部关键点的语句。
        '''
        # 写入左手关键帧
        if isinstance(LeftHandLandmarkList[i], list):
            for j in range(20):
                file.write('\t' + 'at time ' + str(i) + ' $L_' + str(j + 1) + '.pos = [' + str(
                    LeftHandLandmarkList[i][j][0]) + ', ' + str(LeftHandLandmarkList[i][j][2]) + ', ' + str(
                    LeftHandLandmarkList[i][j][1]) + ']\n')

        # 写入右手关键帧
        if isinstance(RightHandLandmarkList[i], list):
            for j in range(20):
                file.write('\t' + 'at time ' + str(i) + ' $R_' + str(j + 1) + '.pos = [' + str(
                    RightHandLandmarkList[i][j][0]) + ', ' + str(RightHandLandmarkList[i][j][2]) + ', ' + str(
                    RightHandLandmarkList[i][j][1]) + ']\n')

    file.write(')\n')
    file.write('animate off\n\n')
    file.close()
    # 此时已生成MS文件，在3ds Max场景中运行该文件即可。

if __name__ == '__main__':
    MC_run()
